/*
 * kdtree.h
 *
 *  Created on: May 1, 2011
 *      Author: tqlong
 */

#ifndef KDTREE_H_
#define KDTREE_H_
#include "tree.h"

template<typename vec_type>
class BoundingBox
{
  vec_type m_min, m_max;
  int n_dim;
public:
  BoundingBox(int dim);

  void accumulate(const vec_type&  v);
  void accumulate(const BoundingBox& b);

  inline const vec_type& min() const { return m_min; }
  inline const vec_type& max() const { return m_max; }
  inline int dim() const { return n_dim; }

  double lowerBound(const vec_type& v, double norm) const;

  void print() const;
};

template<typename DATA>
class KDTree : public Tree<KDTree<DATA>, DATA >
{
public:
  typedef KDTree<DATA>              self_type;
  typedef Tree<self_type,DATA >     __Base;
  typedef DATA                      dataset_type;
  typedef typename DATA::vec_type   vec_type;
  typedef BoundingBox<vec_type>     box_type;
  typedef std::vector<int>          vint_type;
public:
  KDTree(dataset_type& data) : __Base(data), m_box(data.n_dim()) {}
  KDTree(self_type* p, int s, int e) : __Base(p, s, e), m_box(p->n_dim()) {}

  void buildTree();
  void buildStats();
  inline self_type* child(int i) const { return static_cast<self_type*>(this->m_child[i]); }
  self_type* createChild(self_type* p, int s, int e) { return new KDTree(p, s, e); }
  const box_type& box() const { return m_box; }

  void load(const std::string& fName);
protected:
  void split();
  double find_max_variance(int& dim_max) const;
  double find_split_value(int dim) const;
  double min_dim(int dim) const;
  double max_dim(int dim) const;
  void split(int dim, double split_value);

protected:
  box_type m_box;
};

/** An knn visitor example on KDTree */
template <typename T>
struct KNNVisitor
{
  typedef T                        tree_type;
  typedef typename T::dataset_type dataset_type;
  typedef typename T::vec_type     vec_type;
  typedef typename T::box_type     box_type;
  typedef std::vector<int>         vint_type;
  typedef std::vector<double>      vdouble_type;

  vec_type     m_query;
  int K;
  vint_type    m_idx;
  vdouble_type m_dist;

  KNNVisitor(const vec_type& query, int k) :
    m_query(query), K(k), m_idx(k, -1), m_dist(K, POSITIVE_INFINITY)
  {
  }

  bool ok(tree_type* t)
  {
    return t->box().lowerBound(m_query, 2) < m_dist.back();  // use l2 norm
  }

  void process(tree_type* t, int i)
  {
    double d = arma::norm(m_query - t->get(i), 2);           // use l2 norm
    for (int k = 0; k < K; k++) if (d < m_dist[k]) {
      m_idx.insert(m_idx.begin()+k, t->rootIdx(i));
      m_idx.pop_back();
      m_dist.insert(m_dist.begin()+k, d);
      m_dist.pop_back();
      break;
    }
  }

  bool isDone()    // search everything, no stop
  {
    return false;
  }
};

#include "kdtree_impl.h"

#endif /* KDTREE_H_ */
